using System;
using System.Reflection;
using NBody;
using NBody.Analysis;
using NBody.Cosmology;
using OptionParser;

[assembly: AssemblyTitle ("NBodyDensity")]
[assembly: AssemblyVersion ("1.0.*")]
[assembly: AssemblyDescription ("Compute the density in spherical bins.")]
[assembly: AssemblyCopyright ("2005 Joseph D. MacMillan")]

public class DensityOptions : InputOutputGetOptions
{
    [Option ("-b", "Number of bins")]
    public int NumBins;
    
    [Option ("-s", "Bin spacing")]
    public Bin.Spacing Spacing;
    
    [Option ("-c", "Adjust for centre of mass")]
    public bool CentreOfMass;
    
    [Option ("-v", "Limits system to those particles within the virial radius")]
    public bool VirialOnly;
    
    [Option ("-r", "Output relative values: rho / rho_c, r / r_200")]
    public bool Relative;
    
    [Option ("-r2", "Output relative values and multiply density by r^2")]
    public bool Radius2;
    
    [Option ("-cosmo", "Cosmology for background density calculation.")]
    public string CosmoString;
    
    [Option ("-eps", "Softening length")]
    public double Softening = 2.0;
    
    public DensityOptions()
    {
        NumBins = 100;
        Spacing = Bin.Spacing.Logarithmic;
        CosmoString = "SCDM";        
    }
}

public class NBodyDensity
{
    public static double rho_c;
    public static double r_200;
    public static double eps;
    
    public static void Main(string[] args)
    {
        DensityOptions opts = new DensityOptions();
        opts.ProcessArgs(args, ".den");
        eps = opts.Softening;
                        
        NBodySystem s = NBodySystem.Read(opts.InFile);
        
        if (opts.CentreOfMass)
            s.AdjustForCentreOfMass(1e-3);
        
        Bin den = new DensityBin();
        den.BinSpacing = opts.Spacing;
        den.NumBins = opts.NumBins;
        
        if (opts.Relative)
        {
            Model cosmo = Model.FromString(opts.CosmoString);
            rho_c = cosmo.Density(cosmo.Redshift(s.Time));
            r_200 = VirialRadius.CalculateR200(s, rho_c);
            if (opts.VirialOnly)
                s.ExtractSphere(r_200);
            
            den.Analyse(s);
            if (opts.Radius2)
            {
                for (int i = 0; i < den.Data.Length; i++)
                    den.Data[i, 1] *= den.Data[i, 0] * den.Data[i, 0];
                den.Data.Print(opts.OutFile, new DataFile.Transform[] { RelRad, RelDenRad2 });
            }
            else
                den.Data.Print(opts.OutFile, new DataFile.Transform[] { RelRad, RelDen });
        }
        else
        {
            den.Analyse(s);  
            den.Data.Print(opts.OutFile, new DataFile.Transform(Math.Log10));
        }
    }
    
    public static double RelRad(double r)
    {
        // just for fun, return 0 for any radius smaller than the softening lenght
        // (which is hard-coded to be 2 kpc)
        if (r < eps)
            return Math.Log10(0.0);
        else
            return Math.Log10(r / r_200);
    }
    
    public static double RelDen(double rho)
    {
        return Math.Log10(rho / rho_c);
    }
    
    public static double RelDenRad2(double rho)
    {
        return Math.Log10(rho / rho_c / r_200 / r_200);
    }
}
